#!/usr/bin/python3

import gi
gi.require_version("Gtk", "3.0")
from gi.repository import Gtk, Gdk
import os
import argparse

TITLE = "NVIDIA 32-bit Not Found"
TEXT = """<b>32-bit NVIDIA driver files were not found on your host system.</b>
It is recommended that you install them for the best experience.

See <a href="https://github.com/canonical/steam-snap/wiki/FAQ#32-bit-driver">https://github.com/canonical/steam-snap/wiki/FAQ#32-bit-driver</a> for more information."""
COMMAND = """sudo dpkg --add-architecture i386
sudo apt update
sudo apt install libnvidia-gl-{}:i386"""

DO_NOT_SHOW_PATH = os.path.expandvars("$SNAP_USER_COMMON/.nvidia32")

def do_not_show_file(do_not_show: bool):
    """Ok button press.

    `do_not_show`: whether the Do not show again checkbox is checked or not
    """
    if do_not_show:
        open(DO_NOT_SHOW_PATH, "a").close()
    elif os.path.exists(DO_NOT_SHOW_PATH):
        os.remove(DO_NOT_SHOW_PATH)


def nvidia_version():
    """Try to find NVIDIA driver version from modinfo.

    Returns a tuple of the full NVIDIA version triplet and NVIDIA major version.
    """
    modinfo = os.popen(
        "modinfo /usr/lib/modules/$(uname -r)/updates/dkms/nvidia.ko* 2> /dev/null"
        " | grep -m 1 '^version:'"
        " | sed 's/version:\s*//'"
    ).read().splitlines()
    modinfo = modinfo or os.popen(
        "modinfo /usr/lib/modules/$(uname -r)/kernel/nvidia*/nvidia.ko* 2> /dev/null"
        " | grep -m 1 '^version:'"
        " | sed 's/version:\s*//'"
    ).read().splitlines()

    if not modinfo:
        return ()
    nvidia_version = modinfo[0].strip()
    nvidia_version_major = nvidia_version.split(".")[0]
    return (nvidia_version, nvidia_version_major)


def nvidia_missing():
    """Returns whether NVIDIA 32 is missing or not.
    """
    version = nvidia_version()
    if not version:
        return False

    print(f"Found NVIDIA version: {version[0]}")
    needs_32 = os.path.exists(f"/var/lib/snapd/hostfs/usr/lib/x86_64-linux-gnu/libnvidia-glcore.so.{version[0]}")\
        and not os.path.exists(f"/var/lib/snapd/hostfs/usr/lib/i386-linux-gnu/libnvidia-glcore.so.{version[0]}")
    print(f"Need NVIDIA 32-bit: {needs_32}")
    return needs_32


class NvidiaWindow(Gtk.Window):
    def __init__(self):
        super().__init__()
        self.set_resizable(False)
        self.set_border_width(10)
        self.set_position(Gtk.WindowPosition.CENTER)
        self.set_keep_above(True)

        header = Gtk.HeaderBar(title=TITLE)
        self.set_titlebar(header)

        self.grid = Gtk.Grid(
            row_spacing=10,
            column_spacing=10
        )
        self.add(self.grid)

        # main text
        self.label = Gtk.Label()
        self.label.set_markup(TEXT)
        self.grid.attach(self.label, 0, 0, 2, 1)

        # run prompt
        self.run_label = Gtk.Label(
            margin_top=20,
            halign=Gtk.Align.START
        )
        self.run_label.set_markup("To install, run:")
        self.grid.attach_next_to(self.run_label, self.label, Gtk.PositionType.BOTTOM, 1, 1)

        # copy button
        self.copy_btn = Gtk.Button.new_from_icon_name("edit-copy-symbolic", Gtk.IconSize.MENU)
        self.copy_btn.set_margin_top(20)
        self.copy_btn.set_tooltip_text("Copy")
        self.copy_btn.connect("clicked", lambda x: self.button_copy())
        self.copy_btn.set_halign(Gtk.Align.END)
        self.grid.attach_next_to(self.copy_btn, self.run_label, Gtk.PositionType.RIGHT, 1, 1)

        # command text
        self.nvidia_version = nvidia_version()
        self.textview = Gtk.TextView(
            editable=False,
            monospace=True,
            top_margin=5,
            bottom_margin=5,
            right_margin=10,
            left_margin=10
        )
        command_text = COMMAND.format("<your driver version>")
        if self.nvidia_version: command_text = COMMAND.format(self.nvidia_version[1])
        self.textview.get_buffer().set_text(command_text)
        self.grid.attach_next_to(self.textview, self.run_label, Gtk.PositionType.BOTTOM, 2, 1)

        # Do not show again button
        self.checkbox = Gtk.CheckButton(
            label="Do not show again",
            halign=Gtk.Align.END,
            tooltip_markup="This window can always be shown again by running "
                "<tt>$SNAP/bin/nvidia32 --reset</tt> in the Steam Snap shell."
        )
        self.grid.attach_next_to(self.checkbox, self.textview, Gtk.PositionType.BOTTOM, 1, 1)

        # Ok button
        self.btn = Gtk.Button(label="Ok")
        self.btn.connect("clicked", lambda x: self.button_ok(self.checkbox.get_active()))
        self.grid.attach_next_to(self.btn, self.checkbox, Gtk.PositionType.RIGHT, 1, 1)


    def button_ok(self, toggled: bool):
        do_not_show_file(toggled)
        self.close()


    def button_copy(self):
        clipboard = Gtk.Clipboard.get(Gdk.SELECTION_CLIPBOARD)
        clipboard.set_text(COMMAND.format(self.nvidia_version[1]), -1)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--reset",
        help="Show the dialog again if you selected 'Do not show again'",
        action="store_true"
    )
    parser.add_argument(
        "--show",
        help="Show the dialog regardless of conditions met",
        action="store_true"
    )
    args = parser.parse_args()

    if args.reset:
        do_not_show_file(False)

    if args.show or (nvidia_missing() and not os.path.exists(DO_NOT_SHOW_PATH)):
        win = NvidiaWindow()
        win.connect("destroy", Gtk.main_quit)
        win.show_all()
        Gtk.main()


if __name__ == "__main__":
    main()
